/**
 * Copyright (c) Huawei Technologies Co., Ltd. 2023. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "fallback_common.h"
#include "fallback_opapi.h"

#ifdef __cplusplus
extern "C" {
#endif

namespace fallback {
using namespace ge;
using namespace gert;
constexpr size_t kffnInputX = 0;
constexpr size_t kffnInputWeight1 = 1;
constexpr size_t kffnInputWeight2 = 2;
constexpr size_t kffnInputExpertTokens = 3;
constexpr size_t kffnInputbias1 = 4;
constexpr size_t kffnInputbias2 = 5;
constexpr size_t kffnInputscale = 6;
constexpr size_t kffnInputoffset = 7;
constexpr size_t kffnInputdeqScale1 = 8;
constexpr size_t kffnInputdeqScale2 = 9;
constexpr size_t kffnInputantiquantScale1 = 10;
constexpr size_t kffnInputantiquantScale2 = 11;
constexpr size_t kffnInputantiquantOffset1= 12;
constexpr size_t kffnInputantiquantOffset2 = 13;
constexpr size_t kffnOutput = 0;

graphStatus FFNExecuteFunc(OpExecuteContext* host_api_ctx)
{
  OP_CHECK(host_api_ctx == nullptr, OP_LOGE("aclnnfallback", "host_api_ctx is null"), return GRAPH_FAILED);

  auto x_ge = host_api_ctx->GetInputTensor(kffnInputX);
  OP_CHECK(x_ge == nullptr, OP_LOGE("aclnnfallback", "x_ge is null"), return GRAPH_FAILED);

  auto weight1_ge = host_api_ctx->GetInputTensor(kffnInputWeight1);
  OP_CHECK(weight1_ge == nullptr, OP_LOGE("aclnnfallback", "weight1_ge is null"), return GRAPH_FAILED);

  auto weight2_ge = host_api_ctx->GetInputTensor(kffnInputWeight2);
  OP_CHECK(weight2_ge == nullptr, OP_LOGE("aclnnfallback", "weight2_ge is null"), return GRAPH_FAILED);

  auto expert_tokens_ge = host_api_ctx->GetOptionalInputTensor(kffnInputExpertTokens);

  auto bias1_ge = host_api_ctx->GetOptionalInputTensor(kffnInputbias1);

  auto bias2_ge = host_api_ctx->GetOptionalInputTensor(kffnInputbias2);

  auto scale_ge = host_api_ctx->GetOptionalInputTensor(kffnInputscale);

  auto offset_ge = host_api_ctx->GetOptionalInputTensor(kffnInputoffset);

  auto deq_scale1_ge = host_api_ctx->GetOptionalInputTensor(kffnInputdeqScale1);

  auto deq_scale2_ge = host_api_ctx->GetOptionalInputTensor(kffnInputdeqScale2);

  auto antiquant_scale1_ge = host_api_ctx->GetOptionalInputTensor(kffnInputantiquantScale1);

  auto antiquant_scale2_ge = host_api_ctx->GetOptionalInputTensor(kffnInputantiquantScale2);

  auto antiquant_offset1_ge = host_api_ctx->GetOptionalInputTensor(kffnInputantiquantOffset1);

  auto antiquant_offset2_ge = host_api_ctx->GetOptionalInputTensor(kffnInputantiquantOffset2);

  auto output_ge = host_api_ctx->GetOutputTensor(kffnOutput);
  
  std::vector<int64_t> shape;
  if (expert_tokens_ge != nullptr) {
    auto etokens_data_ptr = expert_tokens_ge->GetData<int64_t>();
    auto &expert_tokens_shape = expert_tokens_ge->GetStorageShape();
    if (expert_tokens_shape.GetDimNum() > 0) {
      auto etokens_length = expert_tokens_shape.GetDim(0);
      for (int64_t i = 0; i < etokens_length; i++) {
        int64_t etokens_data = etokens_data_ptr[i];
        shape.push_back(etokens_data);
      }
    }
  }
  auto expert_tokens_int_array = ConvertType(shape);
  
  OP_CHECK(output_ge == nullptr, OP_LOGE("aclnnfallback", "output_ge is null"), return GRAPH_FAILED);
  auto attrs = host_api_ctx->GetAttrs();
  OP_CHECK(attrs == nullptr, OP_LOGE("aclnnfallback", "attrs is null"), return GRAPH_FAILED);
  const char* activation_type_ge = attrs->GetAttrPointer<char>(0);
  const int64_t* inner_pricise_ge = attrs->GetAttrPointer<int64_t>(1);
  const bool* tokens_index_flag_ge = attrs->GetAttrPointer<bool>(3);
  // execute opapi
  auto api_ret = EXEC_OPAPI_CMD(aclnnFFNV2, x_ge, weight1_ge, weight2_ge, expert_tokens_int_array, bias1_ge, bias2_ge, 
                                scale_ge, offset_ge, deq_scale1_ge, deq_scale2_ge, antiquant_scale1_ge, 
                                antiquant_scale2_ge, antiquant_offset1_ge, antiquant_offset2_ge, activation_type_ge, 
                                *inner_pricise_ge, *tokens_index_flag_ge, output_ge);
  OP_CHECK(api_ret != GRAPH_SUCCESS, OP_LOGE("aclnnfallback", "api_ret faild:%d", api_ret), return GRAPH_FAILED);

  return GRAPH_SUCCESS;
}

IMPL_OP(FFN).OpExecuteFunc(FFNExecuteFunc).HostInputs({kffnInputExpertTokens});

}  // namespace fallback

#ifdef __cplusplus
}
#endif
